# -*- coding: utf_8 -*-
import json
import os
import subprocess
import sys

import logging

import re
import redis
import time

from sqlalchemy import create_engine, DateTime, String, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer
from sqlalchemy.orm import sessionmaker

engine = create_engine("postgresql://postgres:postgres@reform-postgres/reform")
Base = declarative_base()

Session = sessionmaker(bind=engine)
session = Session()

# this program should be run as 20 processes, each device a process
# this program call the compiler to compiler, then test user's program, then run user's program


if len(sys.argv) < 2:
    print("Too few arguments! Usage: python code_exec.py device(int)")
    exit(1)


def get_logger():
    logger = logging.getLogger("logger")
    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    return logger


logger = get_logger()

reform_dir = '/home/data/reform'
code_dir = '/home/data/reform/code'
exec_dir = '/home/data/reform/exec'
testcase_filename = '/home/data/reform/testcase'
syscall_limit_filename = '/home/data/reform/syscall_limit.cpp'

device = int(sys.argv[1])

# 注意origin_code_filename不应写为os.path.join(code_dir, str(device), str(device), ".cpp")
# 否则结果为 /home/data/reform/code/15/15/.cpp

origin_code_filename = os.path.join(code_dir, str(device), "origin.cpp")
gen_code_filename = os.path.join(code_dir, str(device), "gen.cpp")
exec_filename = os.path.join(exec_dir, str(device), str(device))

redis_pool = redis.ConnectionPool(host='reform-redis', port=6379)
r = redis.Redis(connection_pool=redis_pool)


def get_status_all():

    r = redis.Redis(connection_pool=redis_pool)
    status_read = r.get("reform.status_read")

    if status_read:
        status_read = json.loads(status_read)
    else:
        status_read = [{} for _ in range(0, 20)]

    def get_pwm(device):
        pwm = r.get("reform.pwm.{}".format(device))
        if pwm:
            return int(pwm)
        else:
            return 0

    def get_code_status(device):
        code_status = r.get("reform.code_status.{}".format(device))
        if code_status:
            return code_status
        else:
            return "idle"

    def get_code_status_detail(device):
        code_status_detail = r.get("reform.code_status_detail.{}".format(device))
        if code_status_detail:
            return code_status_detail
        else:
            return ""

    status_all = [{"device": i + 1,
                   "generator_i": status_read[i]["generator_i"] if "generator_i" in status_read[i] else 0,
                   "generator_u": status_read[i]["generator_u"] if "generator_u" in status_read[i] else 0,
                   "payload_i": status_read[i]["payload_i"] if "payload_i" in status_read[i] else 0,
                   "payload_u": status_read[i]["payload_u"] if "payload_u" in status_read[i] else 0,
                   "pwm": get_pwm(i + 1),
                   "code_status": get_code_status(i + 1),
                   "code_status_detail": get_code_status_detail(i + 1)
                   } for i in range(0, 20)]

    return status_all


class User(Base):
    __tablename__ = 'users'

    def __init__(self, id, password, salt):
        self.id = id
        self.password = password
        self.salt = salt

    id = Column(Integer, primary_key=True)
    password = Column(String(64), nullable=False)
    salt = Column(String(32), nullable=False)
    device = Column(Integer, nullable=False, default=0)
    user_group = Column(Integer, nullable=False, default=1)
    register_time = Column(DateTime, nullable=False, server_default=func.now())


class History(Base):
    __tablename__ = 'history'

    def __init__(self, user_id, generator_i, generator_u, payload_i, payload_u, pwm):
        self.t = func.now()
        self.user_id = user_id
        self.generator_i = generator_i
        self.generator_u = generator_u
        self.payload_i = payload_i
        self.payload_u = payload_u
        self.pwm = pwm

    user_id = Column(Integer, nullable=False)
    t = Column(DateTime, nullable=False, server_default=func.now(), primary_key=True)
    generator_i = Column(Integer, nullable=False)
    generator_u = Column(Integer, nullable=False)
    payload_i = Column(Integer, nullable=False)
    payload_u = Column(Integer, nullable=False)
    pwm = Column(Integer, nullable=False)


while True:
    try:
        code_status = r.get("reform.code_status.{}".format(device))
        if not code_status:
            continue
        if code_status == "submitted" or code_status == "compiling":
            logger.info("compiling!")
            r.set("reform.code_status.{}".format(device), "compiling")
            linker = subprocess.Popen("cat {} {} > {}".format(
                    syscall_limit_filename, origin_code_filename, gen_code_filename),
                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True,
                universal_newlines=True)

            compiler = subprocess.Popen(
                "lrun --max-cpu-time 10 --max-real-time 10 "
                "--max-memory 128M --uid 1000 --gid 1000 g++ {} -lseccomp -o {} 3>&1".format(
                    gen_code_filename, exec_filename),
                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True,
                universal_newlines=True)
            line = "".join(compiler.stdout.readlines())
            matchObj = re.search(
                "([\s\S]*)MEMORY\s+(.*)\s+CPUTIME\s+(.*)\s+REALTIME\s+(.*)\s+SIGNALED\s+(.*)\s+EXITCODE\s+(.*)\s+TERMSIG\s+(.*)\s+EXCEED\s+(.*)",
                line, re.M | re.I)
            if matchObj:
                stdout = matchObj.group(1)
                exit_code = matchObj.group(6)
                exceed = matchObj.group(8)
            else:
                logger.error("no matchObj! line:{}".format(line))
                continue

            logger.debug(stdout)
            logger.debug(exit_code)
            logger.debug(exceed)

            if int(exit_code) == 0 and exceed == "none":
                r.set("reform.code_status.{}".format(device), "testing")
                r.set("reform.code_status_detail.{}".format(device), "")
            else:
                r.set("reform.code_status.{}".format(device), "compile_failed")
                if exceed == "REAL_TIME" or exceed == "CPU_TIME":
                    r.set("reform.code_status_detail.{}".format(device), "TLE")
                elif exceed == "MEMORY":
                    r.set("reform.code_status_detail.{}".format(device), "MLE")
                else:
                    r.set("reform.code_status_detail.{}".format(device), "CE")
            continue
        if code_status == "testing":
            logger.info("testing!")
            tester = subprocess.Popen(
                "lrun --max-cpu-time 1.5 --max-real-time 4 "
                "--max-memory 128M --uid 1000 --gid 1000 {}<{} 3>&1".format(
                    exec_filename, testcase_filename),
                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True,
                universal_newlines=True)
            line = "".join(tester.stdout.readlines())

            matchObj = re.search("([\s\S]*)MEMORY\s+(.*)\s+CPUTIME\s+(.*)\s+REALTIME\s+(.*)\s+SIGNALED\s+(.*)\s+EXITCODE\s+(.*)\s+TERMSIG\s+(.*)\s+EXCEED\s+(.*)", line, re.M | re.I)
            if matchObj:
                # stdout: the output
                stdout = matchObj.group(1)
                exit_code = matchObj.group(6)
                exceed = matchObj.group(8)
            else:
                logger.error("no matchObj! maybe testcase missing or exec file missing?")
                continue

            logger.debug(stdout)
            logger.debug(exit_code)
            logger.debug(exceed)

            if int(exit_code) == 0 and exceed == "none":
                try:
                    lines = stdout.splitlines()
                    if len(lines) == 0:
                        # system call limit triggered
                        r.set("reform.code_status.{}".format(device), "test_failed")
                        r.set("reform.code_status_detail.{}".format(device), "RE")
                        continue
                    for case_out in lines:
                        logger.debug(case_out)

                        pwm = int(case_out)
                        if pwm < 0 or pwm > 10000:
                            raise Exception
                except Exception as e:
                    r.set("reform.code_status.{}".format(device), "test_failed")
                    r.set("reform.code_status_detail.{}".format(device), "WA")
                    continue
                r.set("reform.code_status.{}".format(device), "running")
                r.set("reform.code_status_detail.{}".format(device), "")
            else:
                r.set("reform.code_status.{}".format(device), "test_failed")

                if exceed == "REAL_TIME" or exceed == "CPU_TIME":
                    r.set("reform.code_status_detail.{}".format(device), "TLE")
                elif exceed == "MEMORY":
                    r.set("reform.code_status_detail.{}".format(device), "MLE")
                else:
                    r.set("reform.code_status_detail.{}".format(device), "RE")
            continue
        if code_status == "running":
            logger.info("running!")
            status = get_status_all()[device-1]

            runner = subprocess.Popen(
                "lrun --max-cpu-time 1.5 --max-real-time 4 "
                "--max-memory 128M --uid 1000 --gid 1000 {} 3>&1".format(
                    exec_filename),
                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True,
                universal_newlines=True)
            runner.stdin.write("{} {} {} {} ".format(status["generator_i"], status["generator_u"],
                                                    status["payload_i"], status["payload_u"]))
            line = "".join(runner.stdout.readlines())

            matchObj = re.search(
                "([\s\S]*)MEMORY\s+(.*)\s+CPUTIME\s+(.*)\s+REALTIME\s+(.*)\s+SIGNALED\s+(.*)\s+EXITCODE\s+(.*)\s+TERMSIG\s+(.*)\s+EXCEED\s+(.*)",
                line, re.M | re.I)
            if matchObj:
                # stdout: the output
                stdout = matchObj.group(1)
                exit_code = matchObj.group(6)
                exceed = matchObj.group(8)
            else:
                logger.error("no matchObj! maybe testcase missing or exec file missing?")
                continue
            logger.debug(stdout)
            logger.debug(exit_code)
            logger.debug(exceed)

            try:
                pwm = int(stdout)
                if pwm < 0 or pwm > 10000:
                    raise Exception
            except Exception as e:
                r.set("reform.code_status_detail.{}".format(device), stdout[:10])
                continue
            r.set("reform.pwm.{}".format(device), str(pwm))
            user_id = r.get("reform.device_user.{}".format(device))
            if user_id is None:
                user = session.query(User).filter(User.device == device).first()
                if user is not None:
                    user_id = user.id
                    r.set("reform.device_user.{}".format(device), user_id, ex=60)
            history = History(user_id=user_id, generator_i=status["generator_i"], generator_u=status["generator_u"],
                              payload_i=status["payload_i"], payload_u=status["payload_u"], pwm=pwm)
            session.add(history)
            session.commit()
    except Exception as e:
        logger.exception(e)
    finally:
        time.sleep(0.5)
